import time
import pandas as pd
import torch
from torch.utils.throughput_benchmark import format_time
from preproces_data import preprocessor
import warnings
import numpy as np
warnings.filterwarnings("ignore")


# 为训练模型定义一个函数
def train():
    print("\n训练中.....")

    # 将模型设置为训练模式 - Dropout 层被激活
    model.train()

    # 记录当前时间
    t0 = time.time()

    # 将损失和准确率初始化为 0
    total_loss, total_accuracy = 0, 0

    # 创建一个空列表以保存模型的预测结果
    total_preds = []

    # 对于每个 batch
    for step, batch in enumerate(train_dataloader):

        # 每经过 40 个 batch 后更新进度
        if step % 40 == 0 and not step == 0:
            # 计算经过的时间（以分钟为单位）
            elapsed = format_time(time.time() - t0)

            # 报告进度
            print('  Batch {:>5,}  of  {:>5,}.    Elapsed: {:}.'.format(step, len(train_dataloader), elapsed))

        # 将 batch 推送到 GPU 上
        batch = tuple(t.to(device) for t in batch)

        # 解包 batch 为单独的变量
        # `batch` 包含三个 PyTorch 张量：
        #   [0]: 输入 id
        #   [1]: 注意力掩码
        #   [2]: 标签
        sent_id, mask, labels = batch

        # 在执行反向传播之前，始终清除之前计算的梯度。
        # PyTorch 不会自动执行此操作。
        model.zero_grad()

        # 执行前向传播。这将返回模型的预测结果
        preds = model(sent_id, mask)

        # 计算实际值和预测值之间的损失
        loss = cross_entropy(preds, labels)

        # 累积所有 batch 的训练损失，以便在结束时计算平均损失。
        # `loss` 是一个包含单个值的张量；`.item()` 函数只返回张量中的 Python 值。
        total_loss = total_loss + loss.item()

        # 执行反向传播以计算梯度。
        loss.backward()

        # 使用计算出的梯度更新参数。
        # 优化器决定了“更新规则”——参数如何根据梯度、学习率等进行修改。
        optimizer.step()

        # 模型的预测结果存储在 GPU 上。因此，将其推送到 CPU 上
        preds = preds.detach().cpu().numpy()

        # 累积每个 batch 的模型预测结果
        total_preds.append(preds)

    # 计算一个 epoch 的训练损失
    avg_loss = total_loss / len(train_dataloader)

    # 预测结果的形式为 (batch 数量, batch 大小, 类别数量)。
    # 因此，将预测结果重新整形为 (样本数量, 类别数量)
    total_preds = np.concatenate(total_preds, axis=0)

    # 返回损失和预测结果
    return avg_loss, total_preds


# 为评估模型定义一个函数
def evaluate():
    print("\n评估中.....")

    # 将模型设置为训练模式 - Dropout 层被停用
    model.eval()

    # 记录当前时间
    t0 = time.time()

    # 将损失和准确率初始化为 0
    total_loss, total_accuracy = 0, 0

    # 创建一个空列表以保存模型的预测结果
    total_preds = []
    # 创建空字典保存labels计算准确率
    acc_labels = []

    # 对于每个 batch
    for step, batch in enumerate(validation_dataloader):
        # 每经过 40 个 batch 后更新进度
        if step % 40 == 0 and not step == 0:
            # 计算经过的时间（以分钟为单位）
            elapsed = format_time(time.time() - t0)

            # 报告进度
            print('  Batch {:>5,}  of  {:>5,}.    Elapsed: {:}.'.format(step, len(validation_dataloader), elapsed))

        # 将 batch 推送到 GPU 上
        batch = tuple(t.to(device) for t in batch)

        # 解包 batch 为单独的变量
        # `batch` 包含三个 PyTorch 张量：
        #   [0]: 输入 id
        #   [1]: 注意力掩码
        #   [2]: 标签
        sent_id, mask, labels = batch
        # print("labels:",labels.tolist())
        acc_labels.append(labels.tolist())
        # 在执行前向传播时，停用自动求导
        with torch.no_grad():

            # 执行前向传播。这将返回模型的预测结果
            preds = model(sent_id, mask)

            # 计算实际值和预测值之间的验证损失
            loss = cross_entropy(preds, labels)

            # 累积所有 batch 的验证损失，以便在结束时计算平均损失。
            # `loss` 是一个包含单个值的张量；`.item()` 函数只返回张量中的 Python 值
            total_loss = total_loss + loss.item()
            # acc = preds.argmax(1).eq(labels).sum().item()
            # 将模型的预测结果从 GPU 推送到 CPU
            preds = preds.detach().cpu().numpy()


            # 累积每个 batch 的模型预测结果
            total_preds.append(preds)
    # print("acc:",acc/labels.size(0))
    # 计算一个 epoch 的验证损失
    avg_loss = total_loss / len(validation_dataloader)

    # 预测结果的形式为 (batch 数量, batch 大小, 类别数量)。
    # 因此，将预测结果重新整形为 (样本数量, 类别数量)
    total_preds = np.concatenate(total_preds, axis=0)

    return avg_loss, total_preds, acc_labels





df = pd.read_csv('../datasets/train张大仙.csv')
print(df.shape)
print(df['txt'].sample(5))
# print(df['txt'].value_counts(normalize = True))
# class_counts = df['txt'].value_counts().tolist()
# print(class_counts)
# print(df['txt'].to_numpy())
# words = [jieba.lcut(df['txt'].to_numpy()[i]) for i in range(len(df['txt']))]
# print(words)

# 执行文本清洗
df['clean_text'] = df['txt'].apply(preprocessor)

# 将清理后的文本和标签保存到变量中
text = df['clean_text'].values
labels = df['label'].values.astype('int')

# print("labels:", labels)

# 导入标签编码器
from sklearn.preprocessing import LabelEncoder

# 定义标签编码器
le = LabelEncoder()

# 将目标字符串进行拟合和转换为数字
labels = le.fit_transform(labels.astype('int'))

print("classes：", le.classes_)
print("labels：", labels)

import matplotlib.pyplot as plt

num = [i for i in labels]
plt.hist(num, bins=10)
plt.title("Classes Distribution")
plt.show()

# 导入进度条库
from tqdm import *
from transformers import BertModel, BertTokenizer

model_name = '../bert-base-chinese'
model_path = '../bert-base-chinese'
tokenizer = BertTokenizer.from_pretrained(model_name)
bert_model = BertModel.from_pretrained(model_path)
# 创建一个空列表来保存整数序列
sent_id = []

# 遍历每个推文
for i in tqdm(range(len(text))):
    encoded_sent = tokenizer.encode(text[i],
                                    add_special_tokens=True,
                                    max_length=50,
                                    truncation=True,
                                    pad_to_max_length='right')
    # 将整数序列保存到列表中
    sent_id.append(encoded_sent)

# print("sent_id[0]:", sent_id[0])
# print("sent_id[1]:", sent_id[1])
# print("ids_to_tokens[0]:", tokenizer.convert_ids_to_tokens(sent_id[0]))
# print("ids_to_tokens[1]:", tokenizer.convert_ids_to_tokens(sent_id[1]))

# 创建注意力掩码
attention_masks = []
for sent in sent_id:
    att_mask = [int(token_id > 0) for token_id in sent]
    attention_masks.append(att_mask)


# 使用 train_test_split 将数据划分为训练集和验证集
from sklearn.model_selection import train_test_split

# 使用 90% 的数据作为训练集，10% 作为验证集
train_inputs, validation_inputs, train_labels, validation_labels = train_test_split(sent_id, labels, random_state=2018,
                                                                                    test_size=0.2, stratify=labels)

# 同样处理注意力掩码
train_masks, validation_masks, _, _ = train_test_split(attention_masks, labels, random_state=2018, test_size=0.2,
                                                       stratify=labels)
from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler

# 对于在特定任务上微调 BERT，作者建议使用批量大小为 16 或 32。
# 定义批量大小
batch_size = 32

# 创建训练集的 DataLoader。
# 将张量封装成数据集

train_data = TensorDataset(torch.tensor(train_inputs), torch.tensor(train_masks), torch.tensor(train_labels))

# 定义用于对数据进行采样的采样器
# 随机采样器从数据集中随机采样
# 顺序采样器按顺序采样，总是以相同的顺序
'''
通常使用随机抽样方法来训练模型。这是因为在训练过程中，
模型可能会过度拟合，这意味着模型在训练数据上表现得很好，
但在未见过的测试数据上表现不佳。随机抽样可以帮助增加模型的泛化能力，
使其在各种情况下都能表现得更好。
'''
train_sampler = RandomSampler(train_data)

# 创建训练集的迭代器
train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=batch_size)

# 创建验证集的 DataLoader。
# 将张量封装成数据集
validation_data = TensorDataset(torch.tensor(validation_inputs), torch.tensor(validation_masks), torch.tensor(validation_labels))

# 定义顺序采样器
# 这个采样器按顺序采样数据
validation_sampler = SequentialSampler(validation_data)

# 创建验证集的迭代器
validation_dataloader = DataLoader(validation_data, sampler=validation_sampler, batch_size=batch_size)

# 创建一个迭代器对象
iterator = iter(train_dataloader)

# 加载批量数据
sent_id, mask, target = iterator.__next__()
# print(sent_id)


outputs = bert_model(sent_id, attention_mask=mask, return_dict=False)
hidden_states = outputs[0]
CLS_hidden_state = outputs[1]
# print("Shape of Hidden States:", hidden_states.shape)
# print("Shape of CLS Hidden State:", CLS_hidden_state.shape)


for param in bert_model.parameters():
    param.requires_grad = False

import torch.nn as nn


device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 创建模型
from Classifier_model import Classifier
model = Classifier(bert_model)

# 如果可用，将模型推送到 GPU
model = model.to(device)

# 将张量推送到 GPU
sent_id = sent_id.to(device)
mask = mask.to(device)
target = target.to(device)

# 将输入传递给模型
outputs = model(sent_id, mask)
# print(outputs)

# Adam 优化器
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)


import numpy as np
from sklearn.utils.class_weight import compute_class_weight

# 计算类别权重
class_weights = compute_class_weight(class_weight="balanced", classes=np.unique(labels), y=labels)
print("Class Weights:", class_weights)


weights = torch.tensor(class_weights, dtype=torch.float)

# 将权重传输到 GPU
weights = weights.to(device)

# 定义损失函数
cross_entropy = nn.NLLLoss(weight=weights)

# 计算损失
loss = cross_entropy(outputs, target)
# print("Loss:", loss)

# 将初始损失设为无穷大
best_valid_loss = float('inf')
'''
在注意力掩码中，我们的输入是0和1，但是在最终的计算时，会将在将无效位置的注意力权重设置为一个很小的值，通常为负无穷（-inf），以便在计算注意力分数时将其抑制为接近零的概率。
'''


# 创建一个空列表来存储每个 epoch 的训练和验证损失
train_losses = []
valid_losses = []

epochs = 5



# 对于每个 epoch
for epoch in range(epochs):

    print('\n....... 第 {:} / {:} 个周期 .......'.format(epoch + 1, epochs))

    # 训练模型
    train_loss, _ = train()

    # 评估模型
    valid_loss, total_preds ,acc_labels = evaluate()

    # 保存最佳模型
    if valid_loss < best_valid_loss:
        best_valid_loss = valid_loss
        torch.save(model.state_dict(), 'saved_weights.pt')

    # 累积训练和验证损失
    train_losses.append(train_loss)
    valid_losses.append(valid_loss)
    label_acc=0
    acc_totle=0
    print(f'\n训练损失: {train_loss:.3f}')
    print(f'验证损失: {valid_loss:.3f}')
    accuracy=0
    totle = []
    for i in acc_labels:
        for x in i:
            totle.append(x)

    # print(len(total_preds),len(totle))torch.tensor(total_preds)
    accuracy = torch.tensor(total_preds).argmax(1).eq(torch.tensor(totle)).sum().item()
    # if len(total_preds)==len(totle):
    #     for i in range(len(total_preds)):
    #         if total_preds[i]==totle[i]:
    #             label_acc+=1
    #             acc_totle+=1
    #         else:
    #             acc_totle += 1
    print(f'准确率: {accuracy/len(total_preds):.2f}')



print("")
print("训练完成！")



